import numpy as np

from .base import BaseAMM

from market import Market


class DPM(BaseAMM):
    @staticmethod
    def _price(true_share: int, false_share: int) -> float:
        return true_share / np.sqrt(true_share ** 2 + false_share ** 2)

    def probability_after_trade(self, market: Market, is_true: bool, share: int) -> float:
        if is_true:
            return self._price(market.true_share + share, market.false_share)

        return self._price(market.true_share, market.false_share + share) ** 2

    def trade(self, market: Market, is_true: bool, share: float):
        if is_true:
            market.true_share += share
        else:
            market.false_share += share

    def price(self, market: Market) -> float:
        return self._price(market.true_share, market.false_share)

    def probability(self, market: Market) -> float:
        """
        LMSR中，选项价格就是选项概率
        """
        return self.price(market) ** 2
